﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace InformationRetrieval.Wildcard
{
    /// <summary>
    /// Implements a k-gram index, which allows for efficient wildcard queries.
    /// Usage:
    /// 1. Create a new instance of the k-gram index
    /// 2. Add an arbitrary number of terms to the index
    /// 3. Search the index with a wildcard query
    ///     -> it will return all the terms added in step 2 which match the query.
    /// </summary>
    public class KGramIndex : IWildcardIndex
    {
        #region Fields
        private int _k = 3;
        private Dictionary<string, SortedSet<string>> _index = new Dictionary<string, SortedSet<string>>();
        private HashSet<string> _terms = new HashSet<string>(); 
        #endregion

        #region Constructor
        /// <summary>
        /// Creates a new k-gram index with 'k' set to 3
        /// </summary>
        public KGramIndex() : this(3) { }

        /// <summary>
        /// Creates a new k-gram index with 'k' set to the given value
        /// </summary>
        /// <param name="k">Size of each k-gram; must be greater than 1</param>
        public KGramIndex(int k)
        {
            _k = k;
        } 
        #endregion

        #region Add terms to index
        /// <summary>
        /// Adds a term to the index.
        /// </summary>
        /// <param name="term">The term to be added</param>
        public void AddTerm(string term)
        {
            term = term.ToLower();
            if (!_terms.Contains(term))
            {
                _terms.Add(term);
                foreach (string kgram in GenerateKGrams(term))
                {
                    if (_index.ContainsKey(kgram))
                        _index[kgram].Add(term);
                    else _index[kgram] = new SortedSet<string> { term };
                }
            }
        }

        /// <summary>
        /// Adds each term in the sequence to the index.
        /// </summary>
        /// <param name="terms">A sequence of terms</param>
        public void AddTerms(IEnumerable<string> terms)
        {
            foreach (string term in terms)
            {
                AddTerm(term);
            }
        } 
        #endregion

        #region Find terms in index
        /// <summary>
        /// Finds all terms in the k-gram index that match the query. The query
        /// must be a single term, i.e. it can not contain white-space characters.
        /// The query may contain an arbitrary number of wildcards '*'. 
        /// </summary>
        /// <param name="query">Query string</param>
        /// <returns>A sequence of terms that match the query</returns>
        public IEnumerable<string> Find(string query)
        {
            if (query.Contains(' '))
                throw new ArgumentException("Query string must be a single term without white space characters");

            if (query.Contains('*'))
                return FindWithWildcard(query);
            else return new[] { query };
        }

        /// <summary>
        /// Finds all terms in the k-gram index that match the query string. 
        /// The query string may contain an arbitrary number of wildcards '*'. 
        /// </summary>
        /// <param name="query">Query string with wildcards</param>
        /// <returns>A sequence of terms that match the query string</returns>
        private IEnumerable<string> FindWithWildcard(string query)
        {
            // Get the query's k-grams
            List<string> queryKGrams = GenerateKGrams(query).Where(x => !x.Contains('*')).ToList();

            if (queryKGrams.Count > 0)
            {
                // Find all terms associated with any of the query k-grams
                List<SortedSet<string>> postingLists = new List<SortedSet<string>>();
                foreach (string kGram in queryKGrams)
                {
                    if (_index.ContainsKey(kGram))
                        postingLists.Add(_index[kGram]);
                }

                // No results
                if (postingLists.Count == 0)
                    return new string[] { };

                // Intersect the posting lists -> AND 
                postingLists = postingLists.OrderBy(x => x.Count).ToList();
                SortedSet<string> results = new SortedSet<string>(postingLists[0]);
                foreach (SortedSet<string> postingList in postingLists.Skip(1))
                {
                    results.IntersectWith(postingList);
                }

                // Remove false positives, i.e. results that do not match the query but made it into
                // the result list due to aliasing problems between k-gram-length and query-length
                results.RemoveWhere(x => !WildcardTools.MatchWildcard(x, query));

                return results;
            }
            else
            {
                // There are no query-k-grams -> the query is too short (e.g. just '*')
                // -> we need to do a manual, brute-force search over all terms
                return _terms.Where(x => WildcardTools.MatchWildcard(x, query));
            }
        } 
        #endregion

        #region Helper Methods
        /// <summary>
        /// Splits a term into k-grams, including then starting "$.." and the 
        /// ending "..$" k-gram.
        /// </summary>
        /// <param name="term">A single term</param>
        /// <returns>A sequence of k-grams</returns>
        private IEnumerable<string> GenerateKGrams(string term)
        {
            string tempTerm = "$" + term + "$";
            for (int i = 0; i < tempTerm.Length - _k + 1; i++)
                yield return tempTerm.Substring(i, _k);
        }

        public override string ToString()
        {
            StringBuilder sb = new StringBuilder();

            foreach (var entry in _index)
            {
                sb.AppendLine(String.Format("{0} -> {1}", 
                    entry.Key,
                    String.Join(", ", entry.Value)));
            }


            return sb.ToString();
        }
        #endregion

    }
}
